"""
python fujian1_arima_predict.py
"""
import os
import json
import pandas as pd
from statsmodels.tsa.arima.model import ARIMA

# 输入和输出路径
input_dir = "fujian/fujian1/cubic_spline/interpolation_output"
output_dir = "fujian/fujian1/arima/predict_output"
os.makedirs(output_dir, exist_ok=True)

# 遍历每个 JSON 文件
for filename in os.listdir(input_dir):
    if filename.endswith(".json"):
        filepath = os.path.join(input_dir, filename)
        
        # 读取 JSON 数据
        with open(filepath, 'r', encoding='utf-8') as file:
            data = json.load(file)

        # 转换为 DataFrame
        df = pd.DataFrame(data)
        df['date'] = pd.to_datetime(df['date'])
        df.set_index('date', inplace=True)

        # 显式设置频率
        df = df.asfreq('MS')  # 设置为每月开始

        # 使用 ARIMA 模型进行预测
        model = ARIMA(df['inventory'], order=(1, 1, 1))
        model_fit = model.fit()

        # 预测未来三个月
        forecast = model_fit.forecast(steps=3)
        
        # 生成预测日期
        forecast_dates = pd.date_range(start='2023-07-01', periods=3, freq='MS')
        
        # 准备输出数据，包括原始数据和预测数据
        output_data = [
            {"date": date.strftime('%Y-%m-%d'), "inventory": int(round(inv))}
            for date, inv in zip(df.index, df['inventory'])
        ]
        
        output_data.extend([
            {"date": date.strftime('%Y-%m-%d'), "inventory": int(round(inv))}
            for date, inv in zip(forecast_dates, forecast)
        ])

        # 保存到指定目录
        output_filepath = os.path.join(output_dir, f"predicted_category{filename[18:-5]}.json")
        with open(output_filepath, 'w', encoding='utf-8') as output_file:
            json.dump(output_data, output_file, ensure_ascii=False, indent=4)
        
        print(f"Processed and saved: {output_filepath}")
